{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "644ae2df-d0b9-450e-be21-fb41f96ca368",
   "metadata": {},
   "outputs": [],
   "source": [
    "import struct\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.path import Path\n",
    "import matplotlib as mpl\n",
    "import re\n",
    "import scipy\n",
    "from scipy import stats\n",
    "%config InlineBackend.figure_format='retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "058cfdf5-9b4e-47ff-a5ba-12cb3b70e2d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def ccw_sort(p):\n",
    "    p = np.array(p)\n",
    "    mean = np.mean(p,axis=0)\n",
    "    d = p-mean\n",
    "    s = np.arctan2(d[:,0], d[:,1])\n",
    "    return p[np.argsort(s),:]\n",
    "\n",
    "def great_circle_dist_moon(p1,p2):\n",
    "    \"\"\"\n",
    "    p1 = lon, lat \n",
    "    p2 = lon, lat\n",
    "    \"\"\"\n",
    "    r = 1737.0\n",
    "    lon1 = np.deg2rad(p1[0])\n",
    "    lon2 = np.deg2rad(p2[0])\n",
    "    lat1 = np.deg2rad(p1[1])\n",
    "    lat2 = np.deg2rad(p2[1])\n",
    "    \n",
    "    cent_ang = np.arccos(np.sin(lat1)*np.sin(lat2) + np.cos(lat1)*np.cos(lat2)*np.cos(np.abs(lon1-lon2)))\n",
    "    d = r*cent_ang\n",
    "    return d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5931e769-e21f-4554-9e57-f9cf18283df8",
   "metadata": {},
   "source": [
    "# Load Datasets:\n",
    "\n",
    "### Thorium Data\n",
    "https://pds-geosciences.wustl.edu/missions/lunarp/reduced/thoriumhd.txt "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24482ab2-f1a5-4fe6-8cec-ccf1bbe79d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load thorium data\n",
    "dat_th = pd.read_csv('../data/thoriumhd.txt', skiprows = 13, sep=',\\s+', lineterminator='\\n', names=[\"lat_min\",\"lat_max\",\"lon_min\",\"lon_max\",\"ppm\"])\n",
    "\n",
    "ppm = dat_th.ppm.values.reshape([360,720])\n",
    "thlon = dat_th.lon_min.values.reshape([360,720])\n",
    "thlat = dat_th.lat_min.values.reshape([360,720])\n",
    "\n",
    "\n",
    "## transform into xarray Dataset\n",
    "thor = xr.Dataset(dict(th_ppm = ([\"x\",\"y\"],ppm)),coords = dict(lat=([\"x\",\"y\"],thlat),lon=([\"x\",\"y\"],thlon))) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63375d9a-a140-456b-8b12-659394f56c5b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### TSTA Results File:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e1b5c7-35fb-4dc7-92be-3432a5cf4fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_craterdat = pd.read_csv('../results/tsta_20_200.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01667077-62dd-4925-b0c2-f941c7edc248",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Constants and Functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cefd493f-60a5-4412-a12c-bb4d1880e427",
   "metadata": {},
   "outputs": [],
   "source": [
    "th_c = 0.63 #ppm metzger 1973\n",
    "\n",
    "rho_c = 2550.0 #kg/m**3 Wieczorek 2013  bulk density of highlands crust\n",
    "rho_c_km = rho_c*(1000**3) #kg/km**3\n",
    "\n",
    "D_sc = 18 #km (simple complex transition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e61a530-7a97-4f53-9cec-3444e545c428",
   "metadata": {},
   "outputs": [],
   "source": [
    "###scenario 1\n",
    "def kreep_over(dat,quad,th_k,X):\n",
    "    \"\"\"\n",
    "    dat = pandas dataset \n",
    "    quad = subset of pandas dataset\n",
    "    th_k = numpy array KREEP thorium concentration\n",
    "    X = numpy array Thickness of KREEP layer\n",
    "    \"\"\"\n",
    "    XX,tt = np.meshgrid(X,th_k)\n",
    "\n",
    "    data = np.zeros((len(th_k),len(X),len(dat.loc[quad])))\n",
    "    error = np.zeros_like(data)\n",
    "    a = 0\n",
    "    for _,rows in dat.loc[quad].iterrows():\n",
    "        x_temp = np.zeros_like(XX) #array height of small paraboloid \n",
    "        vc_temp = np.zeros_like(XX)\n",
    "        vk_temp = np.zeros_like(XX)\n",
    "        x_temp[np.where(XX < rows['h_exc'])] = rows['h_exc'] - XX[np.where(XX < rows['h_exc'])] #array height of small paraboloid \n",
    "        vc_temp = 5*np.pi*rows['r_at']*x_temp**2/2.6 ##volume of small paraboloid\n",
    "        vk_temp = -vc_temp + rows['vols']\n",
    "        data[:,:,a] = (tt*vk_temp + th_c*vc_temp)/rows['vols']\n",
    "        error[:,:,a] = data[:,:,a] - rows['ejecta_m_t']\n",
    "        a += 1\n",
    "\n",
    "\n",
    "    rms_err = np.sqrt(np.mean(error**2,axis = 2))\n",
    "    \n",
    "    return rms_err\n",
    "#### scenario 2\n",
    "def kreep_under(dat,quad,th_k,X):\n",
    "    \"\"\"\n",
    "    dat = pandas dataset \n",
    "    quad = subset of pandas dataset\n",
    "    th_k = numpy array KREEP thorium concentration\n",
    "    X = Depth to KREEP layer\n",
    "    \"\"\"\n",
    "    XX,tt = np.meshgrid(X,th_k)\n",
    "\n",
    "    data = np.zeros((len(th_k),len(X),len(dat.loc[quad])))\n",
    "    error = np.zeros_like(data)\n",
    "    a = 0\n",
    "    for _,rows in dat.loc[quad].iterrows():\n",
    "        x_temp = np.zeros_like(XX) #array height of small paraboloid \n",
    "        vc_temp = np.zeros_like(XX)\n",
    "        vk_temp = np.zeros_like(XX)\n",
    "        x_temp[np.where(XX < rows['h_exc'])] = rows['h_exc'] - XX[np.where(XX < rows['h_exc'])] #array height of small paraboloid \n",
    "        vk_temp = 5*np.pi*rows['r_at']*x_temp**2/2.6 ##volume of small paraboloid\n",
    "        vc_temp = -vc_temp + rows['vols']\n",
    "        data[:,:,a] = (tt*vk_temp + th_c*vc_temp)/rows['vols']\n",
    "        error[:,:,a] = data[:,:,a] - rows['ejecta_m_t']\n",
    "        a += 1\n",
    "\n",
    "\n",
    "    rms_err = np.sqrt(np.mean(error**2,axis = 2))\n",
    "    \n",
    "    return rms_err"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c45c58c7-4c23-4101-8c05-e5f06fc7ea60",
   "metadata": {},
   "source": [
    "## Analysis of TSTA Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd9bfb26-450f-4aee-858f-095b138f5ff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_craterdat = full_craterdat.loc[full_craterdat['floor_m_t'].notna()& full_craterdat['ejecta_m_t'].notna()] #exclude nan locations\n",
    "\n",
    "\n",
    "##apparent transient crater\n",
    "full_craterdat['r_at'] = (full_craterdat['rc']**(0.921)*(D_sc/2)**(0.079))/(1.32) #adapted from holsapple 1993 with Rat = Rt/1.3\n",
    "\n",
    "#excavation depth\n",
    "full_craterdat['h_exc'] = 1.3*full_craterdat['r_at']/5.0 #melosh pg. 80\n",
    "\n",
    "full_craterdat['diff'] = full_craterdat['ejecta_m_t'] - full_craterdat['floor_m_t'] \n",
    "full_craterdat['vols'] = np.pi*full_craterdat['h_exc']*full_craterdat['r_at']**2/2\n",
    "\n",
    "\n",
    "\n",
    "full_craterdat = full_craterdat.loc[(full_craterdat['ejecta_m_t'] > 0) & (full_craterdat['floor_m_t'] > 0)]\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0176bf3e-1453-4618-8449-54290b5bf607",
   "metadata": {},
   "source": [
    "# Figure 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c723914-e481-4d96-b5e4-bfcd1fbb16e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "### RMS figure example\n",
    "\n",
    "test_lon = 110.0\n",
    "test_lat = -5.0\n",
    "\n",
    "full_craterdat['dist_from_c'] = great_circle_dist_moon([full_craterdat['lon'],full_craterdat['lat']], [test_lon,test_lat]) #distance from chosen center point\n",
    "\n",
    "local = full_craterdat['dist_from_c'] < 400\n",
    "local_r = local & (full_craterdat['diff'] > full_craterdat['back_std'])\n",
    "local_b = local & (full_craterdat['diff'] < -full_craterdat['back_std'])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "max_hr = full_craterdat.loc[local_r]['h_exc'].max()\n",
    "\n",
    "th_k = np.linspace(th_c,12,1000) # KREEP thorium concentration\n",
    "\n",
    "\n",
    "D = np.linspace(0,max_hr,1000) # Thickness of KREEP layer\n",
    "DD,ttr = np.meshgrid(D,th_k)\n",
    "\n",
    "rms_err_r = kreep_over(full_craterdat,local_r,th_k,D)\n",
    "\n",
    "max_ejr = full_craterdat.loc[local_r]['ejecta_m_t'].max()\n",
    "\n",
    "max_ej_idxr = np.argmin(np.abs(th_k - max_ejr))\n",
    "\n",
    "max_ej_errors = rms_err_r[max_ej_idxr,:]\n",
    "max_ej_idx = np.argmin(max_ej_errors)\n",
    "max_thick = D[max_ej_idx]\n",
    "\n",
    "\n",
    "interr = (max_ej_idxr,max_ej_idx)\n",
    "\n",
    "max_hb = full_craterdat.loc[local_b]['h_exc'].max()\n",
    "\n",
    "\n",
    "X = np.linspace(0,max_hb,1000) # Thickness of KREEP layer\n",
    "XX,ttb = np.meshgrid(X,th_k)\n",
    "\n",
    "rms_err_b = kreep_under(full_craterdat,local_b,th_k,X)\n",
    "\n",
    "max_flb = full_craterdat.loc[local_b]['floor_m_t'].max()\n",
    "\n",
    "max_fl_idxb = np.argmin(np.abs(th_k - max_flb))\n",
    "\n",
    "\n",
    "max_fl_errors = rms_err_b[max_fl_idxb,:] ###rms errors of models with that th_k value\n",
    "min_dep_idx = np.argmin(max_fl_errors) ###index of minimum error \n",
    "min_dep = X[min_dep_idx] ###minimum depth of buried layer\n",
    "\n",
    "interb = (max_fl_idxb,min_dep_idx)\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a84f25da-fd80-448f-9481-f57ce2747b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize = [13,6],constrained_layout=True)\n",
    "gs = fig.add_gridspec(1, 3, width_ratios=(1,1,0.05))\n",
    "ax1 = fig.add_subplot(gs[0,0])\n",
    "\n",
    "min_red_conc = np.min(max_ej_errors)\n",
    "red_level = min_red_conc*np.arange(1.,5.1,.50)## need to come up with way to auto label this\n",
    "\n",
    "\n",
    "\n",
    "rms1 = ax1.contourf(DD,ttr,rms_err_r,cmap = 'binary',levels = red_level,extend = 'both')\n",
    "\n",
    "\n",
    "ej = ax1.plot([0,max_hr],[max_ejr,max_ejr], c = 'firebrick', linestyle = '-', linewidth = 3)\n",
    "mej = ax1.plot([DD[interr],DD[interr]],[0,12],c = 'orangered', linestyle = '--',linewidth = 2)\n",
    "\n",
    "\n",
    "ax1.set_xlabel('Th-rich Layer thickness (km)', fontsize = 16)\n",
    "ax1.set_ylabel('$[Th]_K$ (ppm)', fontsize = 16)\n",
    "ax1.set_title('Scenario 1', size = 20)\n",
    "ax1.tick_params(labelsize=14)\n",
    "#ax.legend(loc = 'upper right',fontsize = 16)\n",
    "ax1.legend([ej[0],mej[0]], ['Minimum [Th] Abundance = {:.1f} ppm'.format(max_ejr),\n",
    "                            'X: Maximum Thickness = {:.1f} km'.format(max_thick)],loc = 'upper right',fontsize = 16)\n",
    "\n",
    "yt = ax1.get_yticks() \n",
    "yt=np.append(yt,0.63)\n",
    "\n",
    "ytl=yt.tolist()\n",
    "ytl[-1]=\"0.63\"\n",
    "ax1.set_yticks(yt)\n",
    "ax1.set_yticklabels(ytl)\n",
    "ax1.set_ylim([0.63,12])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax2 = fig.add_subplot(gs[0,1])\n",
    "\n",
    "\n",
    "\n",
    "min_blu_conc = np.min(max_fl_errors)\n",
    "blue_level = min_blu_conc*np.arange(1.,5.1,.50)## need to come up with way to auto label this\n",
    "\n",
    "blue_level_labels = ['Min Error', '$\\pm$50%','$\\pm$100%','$\\pm$150%','$\\pm$200%','$\\pm$250%','$\\pm$300%', '$\\pm$350%','$\\pm$400%']\n",
    "\n",
    "\n",
    "blue_dep_idx = (max_fl_idxb,min_dep_idx)\n",
    "\n",
    "rms2 = ax2.contourf(XX,ttb,rms_err_b,cmap='binary', levels = blue_level, extend = 'both')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "fl = ax2.plot([0,max_hb],[max_flb,max_flb], c = 'navy', linestyle = '-', linewidth = 3)\n",
    "mfl = ax2.plot([XX[interb],XX[interb]],[0,12], color = 'steelblue',linestyle = '--',linewidth = 2)\n",
    "\n",
    "\n",
    "ax2.set_xlabel('Th-rich Layer depth (km)', fontsize = 16)\n",
    "ax2.set_title('Scenario 2', size = 20)\n",
    "ax2.tick_params(labelleft = False,labelsize=14)\n",
    "ax2.legend([fl[0],mfl[0]], ['Minimum [Th] Abundance = {:.1f} ppm'.format(max_flb),\n",
    "                            'X: Minimum Depth = {:.1f} km'.format(min_dep)],loc = 'upper right',fontsize = 16)\n",
    "\n",
    "yt2 = ax2.get_yticks() \n",
    "yt2=np.append(yt2,0.63)\n",
    "\n",
    "ytl2=yt2.tolist()\n",
    "ytl2[-1]=\"0.63\"\n",
    "ax2.set_yticks(yt2)\n",
    "ax2.set_yticklabels(ytl2)\n",
    "ax2.set_ylim([0.63,12])\n",
    "\n",
    "ax3 = fig.add_subplot(gs[0,2])\n",
    "\n",
    "\n",
    "cb2 = plt.colorbar(rms2, cax = ax3)\n",
    "cb2.set_ticklabels(blue_level_labels)\n",
    "cb2.ax.tick_params(labelsize = 14)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741895ff-8733-45f0-b342-cc23a2659a47",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
